import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

np.random.seed(777)
tf.random.set_seed(777)

BATCH_SIZE = 16
cifar2_dir = '../../../../large_data/DL2/_many_files/cifar2_fast/'


def my_read_files(path):
    y = tf.where(tf.strings.regex_full_match(path, r'.*[\\/]automobile[\\/][^\\/]+$'), 0., 1.)
    img = tf.io.read_file(path)
    img = tf.io.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (32, 32)) / 255.
    return (img, y)


ds = tf.data.Dataset.list_files(cifar2_dir + 'train/*/*', seed=777)\
    .map(my_read_files)\
    .batch(BATCH_SIZE)\
    .prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

spr = 5
spc = 8
spn = 0
for bx, by in ds:
    if spn > spr * spc:
        break
    for i, bxi in enumerate(bx):
        spn += 1
        if spn > spr * spc:
            break
        byi = by[i]
        plt.subplot(spr, spc, spn)
        plt.axis('off')
        plt.title(str(byi.numpy().astype(np.int32)))  # ATTENTION bytes.decode for string and this for numbers
        plt.imshow(bxi)
